#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
配置验证模块 - 增强的配置文件验证和错误提示
提供JSON Schema验证、类型检查、范围验证和自动修复建议
"""

import os
import sys
import json
import re
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple, Union
from datetime import datetime
import logging

# 尝试导入统一日志系统，如果失败则使用标准日志
try:
    from config.unified_logger import get_logger, setup_global_logger
    logger = get_logger(__name__)
except ImportError:
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

class ConfigValidationError(Exception):
    """配置验证错误"""
    def __init__(self, message: str, field: str = None, value: Any = None, suggestions: List[str] = None):
        super().__init__(message)
        self.field = field
        self.value = value
        self.suggestions = suggestions or []

class ConfigValidator:
    """配置验证器"""
    
    def __init__(self):
        self.validation_rules = self._load_validation_rules()
        self.validation_errors = []
        self.validation_warnings = []
        self.auto_fix_suggestions = []
        
    def _load_validation_rules(self) -> Dict[str, Any]:
        """加载验证规则"""
        return {
            "deployment_config": {
                "type": "object",
                "required": ["version", "c2_server", "encryption"],
                "properties": {
                    "version": {
                        "type": "string",
                        "pattern": r"^\d+\.\d+(\.\d+)?$",
                        "description": "配置文件版本号"
                    },
                    "c2_server": {
                        "type": "object",
                        "required": ["host", "port"],
                        "properties": {
                            "host": {
                                "type": "string",
                                "pattern": r"^[a-zA-Z0-9.-]+$",
                                "description": "C2服务器主机名或IP地址"
                            },
                            "port": {
                                "type": "integer",
                                "minimum": 1,
                                "maximum": 65535,
                                "description": "C2服务器端口号"
                            },
                            "use_ssl": {
                                "type": "boolean",
                                "description": "是否使用SSL连接"
                            },
                            "backup_hosts": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "备用主机列表"
                            }
                        }
                    },
                    "encryption": {
                        "type": "object",
                        "required": ["crypto_key", "algorithm"],
                        "properties": {
                            "crypto_key": {
                                "type": "string",
                                "minLength": 16,
                                "maxLength": 64,
                                "description": "加密密钥"
                            },
                            "algorithm": {
                                "type": "string",
                                "enum": ["AES-128", "AES-192", "AES-256", "ChaCha20"],
                                "description": "加密算法"
                            },
                            "use_xor_fallback": {
                                "type": "boolean",
                                "description": "是否使用XOR回退加密"
                            }
                        }
                    },
                    "stealth_options": {
                        "type": "object",
                        "properties": {
                            "stealth_mode": {"type": "boolean"},
                            "command_obfuscation": {"type": "boolean"},
                            "process_hiding": {"type": "boolean"},
                            "traffic_masquerading": {"type": "boolean"}
                        }
                    },
                    "persistence_settings": {
                        "type": "object",
                        "properties": {
                            "multi_layer": {"type": "boolean"},
                            "self_healing": {"type": "boolean"},
                            "health_check_interval": {
                                "type": "integer",
                                "minimum": 60,
                                "maximum": 86400,
                                "description": "健康检查间隔（秒）"
                            },
                            "auto_repair": {"type": "boolean"},
                            "update_check_interval": {
                                "type": "integer",
                                "minimum": 300,
                                "maximum": 604800,
                                "description": "更新检查间隔（秒）"
                            }
                        }
                    },
                    "backdoor_types": {
                        "type": "object",
                        "properties": {
                            "ssh_key": {
                                "type": "object",
                                "properties": {
                                    "enabled": {"type": "boolean"},
                                    "key_comment": {"type": "string"},
                                    "hidden_location": {"type": "boolean"}
                                }
                            },
                            "network": {
                                "type": "object",
                                "properties": {
                                    "enabled": {"type": "boolean"},
                                    "ports": {
                                        "type": "array",
                                        "items": {
                                            "type": "integer",
                                            "minimum": 1,
                                            "maximum": 65535
                                        }
                                    },
                                    "protocols": {
                                        "type": "array",
                                        "items": {
                                            "type": "string",
                                            "enum": ["http", "https", "dns", "tcp", "udp"]
                                        }
                                    }
                                }
                            },
                            "system_service": {
                                "type": "object",
                                "properties": {
                                    "enabled": {"type": "boolean"},
                                    "service_name": {
                                        "type": "string",
                                        "pattern": r"^[a-zA-Z0-9_-]+$"
                                    },
                                    "description": {"type": "string"}
                                }
                            },
                            "crontab": {
                                "type": "object",
                                "properties": {
                                    "enabled": {"type": "boolean"},
                                    "schedule": {
                                        "type": "string",
                                        "pattern": r"^(\*|[0-5]?\d|\*\/\d+)\s+(\*|[01]?\d|2[0-3]|\*\/\d+)\s+(\*|[0-2]?\d|3[01]|\*\/\d+)\s+(\*|[0]?\d|1[0-2]|\*\/\d+)\s+(\*|[0-6]|\*\/\d+)$"
                                    },
                                    "command_obfuscated": {"type": "boolean"}
                                }
                            }
                        }
                    }
                }
            }
        }
    
    def validate_config_file(self, config_path: Union[str, Path]) -> Tuple[bool, List[str], List[str]]:
        """验证配置文件"""
        config_path = Path(config_path)
        self.validation_errors = []
        self.validation_warnings = []
        self.auto_fix_suggestions = []
        
        try:
            # 检查文件是否存在
            if not config_path.exists():
                error_msg = f"配置文件不存在: {config_path}"
                self.validation_errors.append(error_msg)
                self.auto_fix_suggestions.append(f"创建配置文件: {config_path}")
                return False, self.validation_errors, self.auto_fix_suggestions
            
            # 检查文件权限
            if not os.access(config_path, os.R_OK):
                error_msg = f"配置文件无读取权限: {config_path}"
                self.validation_errors.append(error_msg)
                self.auto_fix_suggestions.append(f"修改文件权限: chmod 644 {config_path}")
                return False, self.validation_errors, self.auto_fix_suggestions
            
            # 加载JSON文件
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    config_data = json.load(f)
            except json.JSONDecodeError as e:
                error_msg = f"JSON格式错误: {e}"
                self.validation_errors.append(error_msg)
                self.auto_fix_suggestions.append("检查JSON语法，确保所有括号和引号正确匹配")
                return False, self.validation_errors, self.auto_fix_suggestions
            except UnicodeDecodeError as e:
                error_msg = f"文件编码错误: {e}"
                self.validation_errors.append(error_msg)
                self.auto_fix_suggestions.append("将文件转换为UTF-8编码")
                return False, self.validation_errors, self.auto_fix_suggestions
            
            # 验证配置结构
            is_valid = self._validate_config_structure(config_data)
            
            # 验证配置值
            self._validate_config_values(config_data)
            
            # 检查安全性
            self._check_security_settings(config_data)
            
            # 生成修复建议
            self._generate_fix_suggestions(config_data)
            
            logger.info(f"配置文件验证完成: {config_path}")
            logger.info(f"错误数量: {len(self.validation_errors)}")
            logger.info(f"警告数量: {len(self.validation_warnings)}")
            
            return len(self.validation_errors) == 0, self.validation_errors, self.auto_fix_suggestions
            
        except Exception as e:
            error_msg = f"配置验证过程中发生错误: {e}"
            self.validation_errors.append(error_msg)
            logger.error(error_msg)
            return False, self.validation_errors, self.auto_fix_suggestions
    
    def _validate_config_structure(self, config_data: Dict[str, Any]) -> bool:
        """验证配置结构"""
        try:
            # 检查根级别必需字段
            if "deployment_config" not in config_data:
                self.validation_errors.append("缺少根级别字段: deployment_config")
                self.auto_fix_suggestions.append("添加 'deployment_config' 根字段")
                return False
            
            deployment_config = config_data["deployment_config"]
            rules = self.validation_rules["deployment_config"]
            
            # 检查必需字段
            for required_field in rules.get("required", []):
                if required_field not in deployment_config:
                    self.validation_errors.append(f"缺少必需字段: {required_field}")
                    self.auto_fix_suggestions.append(f"添加必需字段: {required_field}")
            
            # 验证字段类型
            for field_name, field_rules in rules.get("properties", {}).items():
                if field_name in deployment_config:
                    self._validate_field_type(field_name, deployment_config[field_name], field_rules)
            
            return len(self.validation_errors) == 0
            
        except Exception as e:
            self.validation_errors.append(f"结构验证错误: {e}")
            return False
    
    def _validate_field_type(self, field_name: str, field_value: Any, field_rules: Dict[str, Any]):
        """验证字段类型"""
        expected_type = field_rules.get("type")
        
        if expected_type == "string" and not isinstance(field_value, str):
            self.validation_errors.append(f"字段 '{field_name}' 应为字符串类型，实际为: {type(field_value).__name__}")
            self.auto_fix_suggestions.append(f"将字段 '{field_name}' 转换为字符串类型")
        
        elif expected_type == "integer" and not isinstance(field_value, int):
            self.validation_errors.append(f"字段 '{field_name}' 应为整数类型，实际为: {type(field_value).__name__}")
            self.auto_fix_suggestions.append(f"将字段 '{field_name}' 转换为整数类型")
        
        elif expected_type == "boolean" and not isinstance(field_value, bool):
            self.validation_errors.append(f"字段 '{field_name}' 应为布尔类型，实际为: {type(field_value).__name__}")
            self.auto_fix_suggestions.append(f"将字段 '{field_name}' 转换为布尔类型 (true/false)")
        
        elif expected_type == "array" and not isinstance(field_value, list):
            self.validation_errors.append(f"字段 '{field_name}' 应为数组类型，实际为: {type(field_value).__name__}")
            self.auto_fix_suggestions.append(f"将字段 '{field_name}' 转换为数组类型")
        
        elif expected_type == "object" and not isinstance(field_value, dict):
            self.validation_errors.append(f"字段 '{field_name}' 应为对象类型，实际为: {type(field_value).__name__}")
            self.auto_fix_suggestions.append(f"将字段 '{field_name}' 转换为对象类型")
    
    def _validate_config_values(self, config_data: Dict[str, Any]):
        """验证配置值"""
        deployment_config = config_data.get("deployment_config", {})
        
        # 验证版本号格式
        version = deployment_config.get("version")
        if version and not re.match(r"^\d+\.\d+(\.\d+)?$", version):
            self.validation_errors.append(f"版本号格式错误: {version}")
            self.auto_fix_suggestions.append("使用正确的版本号格式，如: '2.1' 或 '2.1.0'")
        
        # 验证C2服务器配置
        c2_server = deployment_config.get("c2_server", {})
        if c2_server:
            self._validate_c2_server(c2_server)
        
        # 验证加密配置
        encryption = deployment_config.get("encryption", {})
        if encryption:
            self._validate_encryption(encryption)
        
        # 验证后门类型配置
        backdoor_types = deployment_config.get("backdoor_types", {})
        if backdoor_types:
            self._validate_backdoor_types(backdoor_types)
    
    def _validate_c2_server(self, c2_server: Dict[str, Any]):
        """验证C2服务器配置"""
        # 验证主机名
        host = c2_server.get("host")
        if host and not re.match(r"^[a-zA-Z0-9.-]+$", host):
            self.validation_errors.append(f"C2服务器主机名格式错误: {host}")
            self.auto_fix_suggestions.append("使用有效的主机名或IP地址")
        
        # 验证端口号
        port = c2_server.get("port")
        if port and (not isinstance(port, int) or port < 1 or port > 65535):
            self.validation_errors.append(f"C2服务器端口号无效: {port}")
            self.auto_fix_suggestions.append("使用1-65535范围内的端口号")
        
        # 验证备用主机
        backup_hosts = c2_server.get("backup_hosts", [])
        if backup_hosts and not isinstance(backup_hosts, list):
            self.validation_errors.append("backup_hosts应为数组类型")
            self.auto_fix_suggestions.append("将backup_hosts设置为字符串数组")
    
    def _validate_encryption(self, encryption: Dict[str, Any]):
        """验证加密配置"""
        # 验证加密密钥
        crypto_key = encryption.get("crypto_key")
        if crypto_key:
            if len(crypto_key) < 16:
                self.validation_errors.append(f"加密密钥长度不足: {len(crypto_key)} < 16")
                self.auto_fix_suggestions.append("使用至少16字符的加密密钥")
            elif len(crypto_key) > 64:
                self.validation_warnings.append(f"加密密钥过长: {len(crypto_key)} > 64")
                self.auto_fix_suggestions.append("建议使用64字符以内的加密密钥")
        
        # 验证加密算法
        algorithm = encryption.get("algorithm")
        valid_algorithms = ["AES-128", "AES-192", "AES-256", "ChaCha20"]
        if algorithm and algorithm not in valid_algorithms:
            self.validation_errors.append(f"不支持的加密算法: {algorithm}")
            self.auto_fix_suggestions.append(f"使用支持的加密算法: {', '.join(valid_algorithms)}")
    
    def _validate_backdoor_types(self, backdoor_types: Dict[str, Any]):
        """验证后门类型配置"""
        # 验证网络后门端口
        network = backdoor_types.get("network", {})
        if network and network.get("enabled"):
            ports = network.get("ports", [])
            for port in ports:
                if not isinstance(port, int) or port < 1 or port > 65535:
                    self.validation_errors.append(f"网络后门端口无效: {port}")
                    self.auto_fix_suggestions.append("使用1-65535范围内的端口号")
        
        # 验证crontab调度格式
        crontab = backdoor_types.get("crontab", {})
        if crontab and crontab.get("enabled"):
            schedule = crontab.get("schedule")
            if schedule and not re.match(r"^(\*|[0-5]?\d|\*\/\d+)\s+(\*|[01]?\d|2[0-3]|\*\/\d+)\s+(\*|[0-2]?\d|3[01]|\*\/\d+)\s+(\*|[0]?\d|1[0-2]|\*\/\d+)\s+(\*|[0-6]|\*\/\d+)$", schedule):
                self.validation_errors.append(f"Crontab调度格式错误: {schedule}")
                self.auto_fix_suggestions.append("使用正确的crontab格式，如: '*/30 * * * *'")
    
    def _check_security_settings(self, config_data: Dict[str, Any]):
        """检查安全设置"""
        deployment_config = config_data.get("deployment_config", {})
        
        # 检查是否启用SSL
        c2_server = deployment_config.get("c2_server", {})
        if not c2_server.get("use_ssl", False):
            self.validation_warnings.append("建议启用SSL连接以提高安全性")
            self.auto_fix_suggestions.append("设置 'use_ssl': true")
        
        # 检查是否启用隐蔽模式
        stealth_options = deployment_config.get("stealth_options", {})
        if not stealth_options.get("stealth_mode", False):
            self.validation_warnings.append("建议启用隐蔽模式")
            self.auto_fix_suggestions.append("设置 'stealth_mode': true")
        
        # 检查加密密钥强度
        encryption = deployment_config.get("encryption", {})
        crypto_key = encryption.get("crypto_key", "")
        if crypto_key and len(crypto_key) < 32:
            self.validation_warnings.append("建议使用更强的加密密钥（32字符以上）")
            self.auto_fix_suggestions.append("生成更长的随机加密密钥")
    
    def _generate_fix_suggestions(self, config_data: Dict[str, Any]):
        """生成修复建议"""
        deployment_config = config_data.get("deployment_config", {})
        
        # 检查是否缺少推荐配置
        if "dns_tunnel" not in deployment_config:
            self.auto_fix_suggestions.append("考虑添加DNS隧道配置以提供备用通信渠道")
        
        if "edr_evasion" not in deployment_config:
            self.auto_fix_suggestions.append("考虑添加EDR规避配置以提高隐蔽性")
        
        if "monitoring" not in deployment_config:
            self.auto_fix_suggestions.append("考虑添加监控配置以跟踪系统状态")
    
    def generate_config_template(self, output_path: Union[str, Path]) -> bool:
        """生成配置文件模板"""
        try:
            template = {
                "deployment_config": {
                    "description": "权限维持工具配置文件",
                    "version": "2.1",
                    "created": datetime.now().isoformat(),
                    "c2_server": {
                        "host": "example.com",
                        "port": 443,
                        "use_ssl": True,
                        "backup_hosts": ["backup1.example.com", "backup2.example.com"]
                    },
                    "encryption": {
                        "crypto_key": "your_32_character_encryption_key_here",
                        "algorithm": "AES-256",
                        "use_xor_fallback": True
                    },
                    "stealth_options": {
                        "stealth_mode": True,
                        "command_obfuscation": True,
                        "process_hiding": True,
                        "traffic_masquerading": True
                    },
                    "persistence_settings": {
                        "multi_layer": True,
                        "self_healing": True,
                        "health_check_interval": 600,
                        "auto_repair": True,
                        "update_check_interval": 3600
                    },
                    "backdoor_types": {
                        "ssh_key": {
                            "enabled": True,
                            "key_comment": "maintenance@server",
                            "hidden_location": True
                        },
                        "network": {
                            "enabled": True,
                            "ports": [443, 80, 53],
                            "protocols": ["https", "dns"]
                        },
                        "system_service": {
                            "enabled": True,
                            "service_name": "system-monitor",
                            "description": "System Monitor Service"
                        },
                        "crontab": {
                            "enabled": True,
                            "schedule": "*/30 * * * *",
                            "command_obfuscated": True
                        }
                    }
                }
            }
            
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(template, f, indent=2, ensure_ascii=False)
            
            logger.info(f"配置模板已生成: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"生成配置模板失败: {e}")
            return False
    
    def auto_fix_config(self, config_path: Union[str, Path], backup: bool = True) -> bool:
        """自动修复配置文件"""
        try:
            config_path = Path(config_path)
            
            # 创建备份
            if backup and config_path.exists():
                backup_path = config_path.with_suffix(f".backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}")
                config_path.rename(backup_path)
                logger.info(f"配置文件已备份到: {backup_path}")
            
            # 加载现有配置
            config_data = {}
            if config_path.exists():
                with open(config_path, 'r', encoding='utf-8') as f:
                    config_data = json.load(f)
            
            # 应用自动修复
            fixed_config = self._apply_auto_fixes(config_data)
            
            # 保存修复后的配置
            with open(config_path, 'w', encoding='utf-8') as f:
                json.dump(fixed_config, f, indent=2, ensure_ascii=False)
            
            logger.info(f"配置文件已自动修复: {config_path}")
            return True
            
        except Exception as e:
            logger.error(f"自动修复配置文件失败: {e}")
            return False
    
    def _apply_auto_fixes(self, config_data: Dict[str, Any]) -> Dict[str, Any]:
        """应用自动修复"""
        # 确保有deployment_config根字段
        if "deployment_config" not in config_data:
            config_data["deployment_config"] = {}
        
        deployment_config = config_data["deployment_config"]
        
        # 添加缺失的必需字段
        if "version" not in deployment_config:
            deployment_config["version"] = "2.1"
        
        if "c2_server" not in deployment_config:
            deployment_config["c2_server"] = {
                "host": "example.com",
                "port": 443,
                "use_ssl": True
            }
        
        if "encryption" not in deployment_config:
            deployment_config["encryption"] = {
                "crypto_key": "default_32_character_key_change_me",
                "algorithm": "AES-256",
                "use_xor_fallback": True
            }
        
        # 修复类型错误
        c2_server = deployment_config.get("c2_server", {})
        if "port" in c2_server and not isinstance(c2_server["port"], int):
            try:
                c2_server["port"] = int(c2_server["port"])
            except ValueError:
                c2_server["port"] = 443
        
        return config_data
    
    def print_validation_report(self):
        """打印验证报告"""
        print("\n" + "="*60)
        print("配置文件验证报告")
        print("="*60)
        
        if self.validation_errors:
            print(f"\n❌ 错误 ({len(self.validation_errors)}):")
            for i, error in enumerate(self.validation_errors, 1):
                print(f"  {i}. {error}")
        
        if self.validation_warnings:
            print(f"\n⚠️  警告 ({len(self.validation_warnings)}):")
            for i, warning in enumerate(self.validation_warnings, 1):
                print(f"  {i}. {warning}")
        
        if self.auto_fix_suggestions:
            print(f"\n💡 修复建议 ({len(self.auto_fix_suggestions)}):")
            for i, suggestion in enumerate(self.auto_fix_suggestions, 1):
                print(f"  {i}. {suggestion}")
        
        if not self.validation_errors and not self.validation_warnings:
            print("\n✅ 配置文件验证通过，无错误或警告")
        
        print("="*60)

def main():
    """主函数 - 命令行工具"""
    import argparse
    
    parser = argparse.ArgumentParser(description="配置文件验证工具")
    parser.add_argument("config_file", help="要验证的配置文件路径")
    parser.add_argument("--auto-fix", action="store_true", help="自动修复配置文件")
    parser.add_argument("--generate-template", help="生成配置模板到指定路径")
    parser.add_argument("--no-backup", action="store_true", help="自动修复时不创建备份")
    
    args = parser.parse_args()
    
    validator = ConfigValidator()
    
    # 生成模板
    if args.generate_template:
        if validator.generate_config_template(args.generate_template):
            print(f"✅ 配置模板已生成: {args.generate_template}")
        else:
            print("❌ 生成配置模板失败")
        return
    
    # 验证配置文件
    is_valid, errors, suggestions = validator.validate_config_file(args.config_file)
    
    # 打印验证报告
    validator.print_validation_report()
    
    # 自动修复
    if args.auto_fix and not is_valid:
        print("\n🔧 正在自动修复配置文件...")
        if validator.auto_fix_config(args.config_file, backup=not args.no_backup):
            print("✅ 配置文件已自动修复")
            # 重新验证
            is_valid, _, _ = validator.validate_config_file(args.config_file)
            if is_valid:
                print("✅ 修复后验证通过")
            else:
                print("⚠️  修复后仍有问题，请手动检查")
        else:
            print("❌ 自动修复失败")
    
    sys.exit(0 if is_valid else 1)

if __name__ == "__main__":
    main()